package com.test.demo.config;

import com.alibaba.druid.pool.DruidDataSource;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import javax.sql.DataSource;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

/**
 * 数据源切换处理
 */
public class DynamicDataSourceContextHolder {
    private static final Logger log = LoggerFactory.getLogger(DynamicDataSourceContextHolder.class);

    /**
     * 存储已创建的数据源
     */
    public static Map<String, DataSource> dataSourceMap = new ConcurrentHashMap<>();

    /**
     * 使用ThreadLocal维护变量，ThreadLocal为每个使用该变量的线程提供独立的变量副本，
     * 所以每一个线程都可以独立地改变自己的副本，而不会影响其它线程所对应的副本。
     */
    private static final ThreadLocal<String> CONTEXT_HOLDER = new ThreadLocal<>();

    /**
     * 设置数据源
     */
    public static void setDataSourceKey(String key) {
        log.info("切换到{}数据源", key);
        CONTEXT_HOLDER.set(key);
    }

    /**
     * 获取数据源
     */
    public static String getDataSourceKey() {
        return CONTEXT_HOLDER.get();
    }

    /**
     * 清除数据源
     */
    public static void clearDataSourceKey() {
        CONTEXT_HOLDER.remove();
    }

    /**
     * 添加数据源
     */
    public static void addDataSource(String key, DataSource dataSource) {
        dataSourceMap.put(key, dataSource);
    }

    /**
     * 删除数据源
     */
    public static void removeDataSource(String key) {
        dataSourceMap.remove(key);
    }

    /**
     * 获取所有数据源
     */
    public static Map<String, DataSource> getDataSourceMap() {
        return dataSourceMap;
    }

    /**
     * 检查数据源是否存在
     */
    public static boolean containsDataSource(String key) {
        return dataSourceMap.containsKey(key);
    }

    /**
     * 关闭数据源
     */
    public static void closeDataSource(String key) {
        if (containsDataSource(key)) {
            DataSource dataSource = dataSourceMap.get(key);
            if (dataSource instanceof DruidDataSource) {
                ((DruidDataSource) dataSource).close();
            }
            removeDataSource(key);
        }
    }
} 